import { NextRequest, NextResponse } from "next/server";
import { Message as VercelChatMessage, StreamingTextResponse } from "ai";

import { ChatOpenAI } from "@langchain/openai";
import { PromptTemplate } from "@langchain/core/prompts";
import { HttpResponseOutputParser } from "langchain/output_parsers";

export const runtime = "edge";

const formatMessage = (message: VercelChatMessage) => {
    return `${message.role}: ${message.content}`;
};

const TEMPLATE = `You are an expert programmer,Chinese answers are the main focus.

Current conversation:
{chat_history}

User: {input}
AI:`;


export async function POST(req: NextRequest) {
    try {
        const body = await req.json();
        const messages = body.messages ?? [];
        const formattedPreviousMessages = messages.slice(0, -1).map(formatMessage);
        const currentMessageContent = messages[messages.length - 1].content;
        const prompt = PromptTemplate.fromTemplate(TEMPLATE);

        const model = new ChatOpenAI({
            temperature: 0,
            modelName: "gpt-3.5-turbo"
        });
        const outputParser = new HttpResponseOutputParser();
        const chain = prompt.pipe(model).pipe(outputParser);

        const result = await chain.stream({
            chat_history: formattedPreviousMessages.join("\n"),
            input: currentMessageContent,
        });
        return new StreamingTextResponse(result);
    } catch (e: any) {
        return NextResponse.json({ error: e.message }, { status: e.status ?? 500 });
    }
}
